"""
用于简单测试单照片的脚本
"""
import os

import cv2

from ultralytics import YOLO


# yolo = YOLO(R"F:\Project\DeepLearning\YOLOv8\material\out\20240108_300_train4\weights\best.pt")
yolo = YOLO("best.pt")
folderPath = R"F:\Project\temp\1"

def run():
    # 关于配置参数,例如此处的task,参考文件在ultralytics/ultralytics/cfg/default.yaml
    # yolo = YOLO("best.pt")
    # yolo = YOLO("../material/yolov8n.pt")

    useOpenCV = False
    if not useOpenCV:
        # 检测并获得结果集
        # yolo.predict(source="../material/bus.jpg", show=True)
        yolo.predict(source="../material/tank0.jpg", show=True)
    else:
        results = yolo.predict(source="../material/test2.jpg")
        # results = yolo.predict(source="../material/test2_gray.jpg")
        # 使用openCV显示结果
        img = results[0].plot()
        img = cv2.resize(img, [1600, 900])
        cv2.imshow("a", img)
    cv2.waitKey()
    cv2.destroyWindow("a")


def runFolder():
    allFiles = os.listdir(folderPath)
    # for root, ds, fs in os.walk(thePath):
    # root当前目录,ds目录下的文件夹,fs是文件名集合
    for fs in allFiles:
        name, suffix = os.path.splitext(fs)
        if suffix == ".jpg" or suffix == ".png":
            print(name, suffix)
            results = yolo.predict(source=folderPath + "/" + fs)
            # 使用openCV显示结果
            img = results[0].plot()
            # width>1600
            if img.shape[1] > 1600:
                scale_percent = 60  # percent of original size
                width = int(img.shape[1] * scale_percent / 100)
                height = int(img.shape[0] * scale_percent / 100)
                img = cv2.resize(img, dsize=(width, height))
            cv2.imshow("a", img)
            cv2.waitKey()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    runFolder()
    # run()
